import numpy as np
import matplotlib.pyplot as plt

# 设置matplotlib的字体为支持中文的字体
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 指定默认字体为微软雅黑
plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像时负号'-'显示为方块的问题

# 5x5的网格代表不同的地区
grid_size = 5
I_grid = np.zeros((grid_size, grid_size))

N = 10000  # 总人口
beta = 0.2  # 感染率
gamma = 0.1  # 疫苗接种率
sigma = 0.01  # 某类感染/恢复/死亡的随机波动率
delta = 0.1  # 暴露者转化为感染者的速率
alpha = 0.05  # 感染者恢复率
rho = 0.01  # 感染者死亡率
dt = 0.1  # 时间步长
T = 10    # 总时间（以dt为单位）

# 初始化
S0, V0, E0, I0, R0, D0 = N - 1, 0, 0, 1, 0, 0  # 初始条件
S, V, E, I, R, D = [S0], [V0], [E0], [I0], [R0], [D0]
np.random.seed(0)

# Euler-Maruyama方法迭代更新
for t in np.arange(dt, T, dt):
    dW_S, dW_V, dW_E, dW_I, dW_R, dW_D = np.random.normal(0, np.sqrt(dt), 6)
    S_new = S[-1] + (-beta * S[-1] * (I[-1] + E[-1]) / N - gamma * S[-1]) * dt + sigma * S[-1] * dW_S
    V_new = V[-1] + (gamma * S[-1] - sigma * V[-1] * (I[-1] + E[-1]) / N) * dt + sigma * V[-1] * dW_V
    E_new = E[-1] + (
                beta * S[-1] * (I[-1] + E[-1]) / N + sigma * V[-1] * (I[-1] + E[-1]) / N - delta * E[-1]) * dt + sigma * \
            E[-1] * dW_E
    I_new = I[-1] + (delta * E[-1] - (alpha + rho) * I[-1]) * dt + sigma * I[-1] * dW_I
    R_new = R[-1] + alpha * I[-1] * dt + sigma * R[-1] * dW_R
    D_new = D[-1] + rho * I[-1] * dt + sigma * D[-1] * dW_D
    S.append(S_new)
    V.append(V_new)
    E.append(E_new)
    I.append(I_new)
    R.append(R_new)
    D.append(D_new)


initial_infected_indices = [(1, 1), (3, 3)]
for i, j in initial_infected_indices:
    I_grid[i, j] = 10

for t in range(int(T // dt)):

    for i in range(grid_size):
        for j in range(grid_size):
            if I_grid[i, j] > 0:
                for di, dj in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                    ni, nj = i + di, j + dj
                    if 0 <= ni < grid_size and 0 <= nj < grid_size:
                        I_grid[ni, nj] += 1

    # 可视化当前时间步长的疫情热力图
    if (t % 10 == 0) or (t == T // dt - 1):
        plt.figure(figsize=(8, 6))
        plt.imshow(I_grid, cmap='hot', interpolation='nearest')  # 使用hot色图显示传染者数量
        plt.colorbar(label='传染者数量')
        plt.title(f'时间步长 {t * dt}')
        plt.xlabel('地区')
        plt.ylabel('地区')
        plt.xticks(range(grid_size), [f'地区{i + 1}' for i in range(grid_size)])
        plt.yticks(range(grid_size), [f'地区{i + 1}' for i in range(grid_size)])
        plt.grid(False)
        plt.show()
